import random
import inspect
from enum import Enum, unique
import unittest


"""
给各个子集的参数加密的类
"""
class FederalStudyEncryptor:
    """
    此枚举类用于表示各个机构类型：未知类型、子集类型、联邦计算平均值机构类型
    """
    @unique
    class TIdentity(Enum):
        kUnknown = 1,
        kSubset = 2,
        kFederalCalculator = 3,

    def __init__(self, subset_cnt: int):
        if subset_cnt <= 0:
            raise AssertionError('subset cnt must be a positive number.')

        self.__subset_cnt = subset_cnt
        self.__subsets = []
        self.__federalCalculator = None

        self.__avg_noise = 0
        self.__noises = []

    def genNoises(self):
        total_noise = 0
        self.__noises.clear()

        for i in range(self.__subset_cnt):
            # 为每个子集随机生成一个-10到10的噪音
            noise = random.randint(-10, 10)
            self.__noises.append(noise)
            total_noise += noise
            # print('%d' % noise, end=' ')
        self.__avg_noise = total_noise / self.__subset_cnt
        # print('avg: %f' % self.__avg_noise)

    def getNoise(self, identity) -> int:
        # 噪音只对对应的子集开放
        self.__authIdentity(FederalStudyEncryptor.TIdentity.kSubset, identity)

        if self.__noises is None or len(self.__noises) == 0:
            raise AssertionError('please call genNoise() first to generate noise.')

        subset_idx = self.__subsets.index(identity)
        return self.__noises[subset_idx]

    def getAvgNoise(self, identity) -> int:
        # 噪音只对对应的子集开放，不对其他机构 比如平均值计算机构开放。
        self.__authIdentity(FederalStudyEncryptor.TIdentity.kSubset, identity)
        return self.__avg_noise

    def registerSubset(self, subset):
        if len(self.__subsets) >= self.__subset_cnt:
            raise RuntimeWarning('only %d subsets are allowed.' % self.__subset_cnt)
        self.__subsets.append(subset)

    def registerFederalCalculator(self, calculator):
        if self.__federalCalculator is not None:
            raise RuntimeWarning('Federal Calculator already registered, ignored.')
        self.__federalCalculator = calculator

    # 鉴定调用者身份，某些数据只对指定机构开放
    def __authIdentity(self, expect, to_auth):
        if expect == FederalStudyEncryptor.TIdentity.kSubset:
            if to_auth in self.__subsets:
                return
        elif expect == FederalStudyEncryptor.TIdentity.kFederalCalculator:
            if to_auth == self.__federalCalculator:
                return
        raise RuntimeWarning('you have NO authority to call <{}()>'.format(inspect.stack()[1][3]))



class Subset:
    def __init__(self):
        pass

    def printNoise(self):
        print(encryptor.getNoise(self))

    def test(self):
        print(encryptor.getAvgNoise(self))



class UnitTestEncryptor(unittest.TestCase):
    def test_authority(self):
        with self.assertRaises(RuntimeWarning):
            encryptor.genNoises()
            subset1 = Subset()
            subset2 = Subset()
            subset3 = Subset()
            encryptor.registerSubset(subset1)
            encryptor.registerSubset(subset2)
            # encryptor.registerSubset(subset3)

            subset1.printNoise()
            subset2.printNoise()
            subset3.printNoise()


encryptor = FederalStudyEncryptor(3)

# 这个地方是用来测试这个类的正确性用的，只用于测试罢了。
if __name__ == '__main__':
    unittest.main()
